#!/bin/bash

set -Eeuo pipefail

DATA=/data

declare -A MOUNTS
MOUNTS["/root/.cache"]="${DATA}/.cache"
MOUNTS["${WORKSPACE}/models"]="${DATA}/models"
MOUNTS["${WORKSPACE}/loras"]="${DATA}/loras"
MOUNTS["${WORKSPACE}/training"]="${DATA}/training"
MOUNTS["${WORKSPACE}/extensions"]="${DATA}/config/oobabooga/extensions"
MOUNTS["${WORKSPACE}/characters"]="${DATA}/config/oobabooga/characters"
MOUNTS["${WORKSPACE}/instruction-templates"]="${DATA}/config/oobabooga/instruction-templates"
MOUNTS["${WORKSPACE}/prompts"]="${DATA}/config/oobabooga/prompts"

for to_path in "${!MOUNTS[@]}"; do
  set -Eeuo pipefail
  # 源目录不存在就初始化源目录
  from_path="${MOUNTS[${to_path}]}"
  if [ ! -d "$from_path" ]; then
    mkdir -vp "$(dirname "$from_path")"
    cp -r "${to_path}" "${from_path}"
    chown -R 1000:1000 "${from_path}"
  fi
  # 目标目录不是符号链接则创建符号链接
  if [ ! -L "${to_path}" ]; then
    rm -rf "${to_path}"
    mkdir -vp "$(dirname "${to_path}")"
    ln -sT "${from_path}" "${to_path}"
  fi
  echo "Mounted $(basename "${from_path}")"
done

# TODO: move all mkdir -p ?
if [ ! -d "/data/config/oobabooga/scripts/" ]; then
  mkdir -p /data/config/oobabooga/scripts/
fi

if [ ! -d "${WORKSPACE}/scripts/" ]; then
  mkdir -p "${WORKSPACE}/scripts/"
fi

# mount scripts individually
# find "${WORKSPACE}/scripts/" -maxdepth 1 -type l -delete
# cp -vrfTs "${DATA}/config/oobabooga/scripts/" "${WORKSPACE}/scripts/"

# mount settings individually
if [ -f "${DATA}/config/oobabooga/settings.yaml" ]; then
  rm "${WORKSPACE}/settings.yaml" || exit 0
  cp "${DATA}/config/oobabooga/settings.yaml" "${WORKSPACE}/settings.yaml"
fi

echo "Installing extension dependencies (if any)"
shopt -s nullglob

# Duplicate code.
# list=(./extensions/*/requirements.txt)
# for requirement in "${list[@]}"; do
#     pip install -r "$requirement"
# done

if [ -f "/data/config/oobabooga/startup.sh" ]; then
  pushd "${WORKSPACE}"
  echo "Running startup script"
  . /data/config/oobabooga/startup.sh
  popd
fi

exec "$@"
